import os
import paddle
import numpy as np
from paddle.vision.transforms import Compose, Resize, ToTensor
from PIL import Image
from utils.dataset import LaneDataset
from model.deeplabv3p import DeepLabv3p
# from model.unet import ResNetUNet
from utils.process_labels import decode_color_labels
from config import Config

device_id = 0
predict_net = 'deeplabv3p'
nets = {'deeplabv3p' : DeepLabv3p}


def concat_and_save_image(color_pred, color_mask, predict_result_path):
    color_pred = Image.fromarray(color_pred, 'RGB').convert('L')
    color_mask = Image.fromarray(color_mask, 'RGB').convert('L')
    height = color_pred.height
    width = color_pred.width
    print('height', height)
    print('width', width)
    target = Image.new('L', (width, height * 2))  #拼接前需要写拼接完成后的图片大小
    # a = size * i         # 图片距离左边的大小
    # b = 0                # 图片距离上边的大小
    # c = size * (i + 1)   # 图片距离左边的大小 + 图片自身宽度
    # d = size             # 图片距离上边的大小 + 图片自身高度
    target.paste(color_pred, box=(0, 0, width, height))
    target.paste(color_mask, box=(0, height, width, height * 2))
    print('拼接图片的路径为：', predict_result_path)
    target.save(predict_result_path)


def load_model(model_path):

    lane_config = Config()
    net = nets[predict_net](lane_config)
    net.eval()

    layer_state_dict = paddle.load(model_path)
    net.set_state_dict(layer_state_dict)

    return net

def img_transaction(img):
    # 将img中Shape中尺寸为1的维度删除，即删除batch的维度
    img = paddle.squeeze(img)
    img = img.numpy()
    print('img shape', img.shape)
    img = decode_color_labels(img)
    # 必须把第一个通道维度放到最后，这样后面PIL模块转换图片才能正常
    img = np.transpose(img, (2, 1, 0))
    return img

def get_color_mask(pred):
    pred = paddle.nn.functional.softmax(pred, axis=1)
    # pred_heatmap = troch.max(pred, dim=1)
    #
    pred = paddle.argmax(pred, axis=1)
    return img_transaction(pred)

def main():
    lane_config = Config()
    model_path = os.path.join(os.getcwd(), lane_config.SAVE_PATH, 'finaNet.pdparams')
    predict_result_path = os.path.join(os.getcwd(), 'predict_result')
    if not os.path.exists(predict_result_path):
        os.mkdir(predict_result_path)
    net = load_model(model_path)

    test_dataset = LaneDataset("test.csv", transform=Compose([Resize(size=(1024, 384)), ToTensor()]))
    test_data_batch = paddle.io.DataLoader(test_dataset, batch_size=1, shuffle=False, drop_last=False)
    i = 0
    for batch_item in test_data_batch:
        image, mask = batch_item[0], batch_item[1]
        pred = net(image)
        color_pred = get_color_mask(pred)
        color_mask = img_transaction(mask)
        image_result_path = os.path.join(predict_result_path, 'result{}.jpg'.format(i))
        concat_and_save_image(color_pred, color_mask, image_result_path)
        break
        i += 1

if __name__ == '__main__':
    main()